# -*- coding: utf-8 -*-
# @Time : 2022/06/07/ 14:42
import ipaddress
import os
import struct
import socket
from urllib.parse import urlparse
import zipfile
import requests




class IPFind:
    _ipv4_instance_location = None
    _ipv4_instance_scenes = None
    _ipv4_instance_proxy = None
    _ipv6_instance = None

    def __init__(self, file_name: str, version: int, read_mode="mem"):
        if not hasattr(self, '_initialized'):
            raise Exception("You must use the factory method to instantiate this object.") # 抛出异常避免实例化
        self.numbers = None
        self.data = None
        self._handle = None
        self.buchang = 9
        self.prefArr = None
        self.version = version
        self.file_name = file_name  # mem:内存版 file:文件版
        self.read_mode = read_mode
        if version == 4:
            self.get_ipv4_object()
        elif version == 6:
            self.get_ipv6_object()


    @classmethod
    def ipv4_instance_location_file(cls, file_name):
        # 使用工厂方法实例化
        instance_location = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
        instance_location._initialized = True  # 标记该对象已初始化
        instance_location.__init__(file_name,4,"file")  # 调用 __init__ 初始化
        return instance_location
    
    @classmethod
    def ipv4_instance_location_memory(cls, file_name):
        # 使用工厂方法实例化
        if not cls._ipv4_instance_location:
            cls._ipv4_instance_location = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
            cls._ipv4_instance_location._initialized = True  # 标记该对象已初始化
            cls._ipv4_instance_location.__init__(file_name,4,"mem")  # 调用 __init__ 初始化
        return cls._ipv4_instance_location

    @classmethod
    def ipv4_instance_scenes_file(cls, file_name):
        # 使用工厂方法实例化
        instance_scenes = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
        instance_scenes._initialized = True  # 标记该对象已初始化
        instance_scenes.__init__(file_name,4,"file")  # 调用 __init__ 初始化
        return instance_scenes
    
    
    @classmethod
    def ipv4_instance_scenes_memory(cls, file_name):
        # 使用工厂方法实例化
        if not cls._ipv4_instance_scenes:
            cls._ipv4_instance_scenes = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
            cls._ipv4_instance_scenes._initialized = True  # 标记该对象已初始化
            cls._ipv4_instance_scenes.__init__(file_name,4,"mem")  # 调用 __init__ 初始化
        return cls._ipv4_instance_scenes
    
    @classmethod
    def ipv4_instance_proxy_file(cls, file_name):
        # 使用工厂方法实例化
        instance_proxy = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
        instance_proxy._initialized = True  # 标记该对象已初始化
        instance_proxy.__init__(file_name,4,"file")  # 调用 __init__ 初始化
        return instance_proxy
    
    
    @classmethod
    def ipv4_instance_proxy_memory(cls, file_name):
        # 使用工厂方法实例化
        if not cls._ipv4_instance_proxy:
            cls._ipv4_instance_proxy = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
            cls._ipv4_instance_proxy._initialized = True  # 标记该对象已初始化
            cls._ipv4_instance_proxy.__init__(file_name,4,"mem")  # 调用 __init__ 初始化
        return cls._ipv4_instance_proxy
    
    @classmethod
    def ipv6_instance_file(cls, file_name):
        # 使用工厂方法实例化
        instance = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
        instance._initialized = True  # 标记该对象已初始化
        instance.__init__(file_name,6,"file")  # 调用 __init__ 初始化
        return cls._ipv6_instance
    
    
    @classmethod
    def ipv6_instance_memory(cls, file_name):
        # 使用工厂方法实例化
        if not cls._ipv6_instance:
            cls._ipv6_instance = super(IPFind,cls).__new__(cls)  # 创建一个新的对象实例
            cls._ipv6_instance._initialized = True  # 标记该对象已初始化
            cls._ipv6_instance.__init__(file_name,6,"mem")  # 调用 __init__ 初始化
        return cls._ipv6_instance
    
    
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.close()

    def close(self):
        self._handle.close()  

    def get_ipv4_object(self):
        self.version = 4
        self._handle = open(self.file_name, "rb")
        if self.read_mode == "mem":
            self.data = self._handle.read()
        self.prefArr = []
        i = 0
        while i < 256:
            p = i * 8 + 4
            self.prefArr.append([self.unpack_int_4byte(p), self.unpack_int_4byte(p + 4)])
            i += 1

    def get_ipv6_object(self):
        self.version = 6
        self._handle = open(self.file_name, "rb")
        if self.read_mode == "mem":
            self.data = self._handle.read()
        self.prefArr = {}
        self.numbers = self.unpack_int_4byte(4)  # IP起始 去重后的数量
        i = 0
        # 还原第一段第二段
        while i < self.numbers:
            p = i * 12 + 4 + 4  # 12 IP起始行所在的行 + IP起始行结束行 + IP起始 的第一位
            self.prefArr[self.unpack_int_4byte(p + 4 + 4)] = [self.unpack_int_4byte(p),
                                                              self.unpack_int_4byte(p + 4)]
            i += 1

    def get(self, ip: str) -> str:
        if self.version == 4:
            ipdot = ip.split('.')
            prefix = int(ipdot[0])
            if prefix < 0 or prefix > 255 or len(ipdot) != 4:
                raise ValueError("invalid ip address")
            intIP = self.ip_to_int(ip)
            low = self.prefArr[prefix][0]
            high = self.prefArr[prefix][1]
            cur = low if low == high else self.ipv4_search(low, high, intIP)
            return self.get_ipv4_addr(cur)
        elif self.version == 6:
            ipdot = ip.split(':')
            prefix = int(ipdot[0], 16)
            # intIP = self.ip_to_int(ip)
            try:
                intIP = int(ipaddress.IPv6Address(ip))
            except:
                raise ValueError("invalid ip address")
            # print(intIP)
            if prefix not in self.prefArr:
                return "no information"
            low = self.prefArr[prefix][0]
            high = self.prefArr[prefix][1]
            cur = low if low == high else self.ipv6_search(low, high, intIP)
            return self.get_ipv6_addr(cur)

    def ipv4_search(self, low: int, high: int, k: int) -> int:
        M = 0
        while low <= high:
            mid = (low + high) // 2
            end_ip_num = self.unpack_int_4byte(2052 + (mid * self.buchang))
            if end_ip_num >= k:
                M = mid
                if mid == 0:
                    break
                high = mid - 1
            else:
                low = mid + 1
        return M

    def ipv6_search(self, low, high, k):
        M = 0
        while low <= high:
            mid = (low + high) // 2
            p = self.numbers * 12 + 4 + 4 + (mid * 55)
            if self.read_mode == "mem":
                end_ip_num = int(self.data[p:p + 50].decode('utf-8').replace("*", ""))
            elif self.read_mode == "file":
                self._handle.seek(p)
                end_ip_num = int(self._handle.read(50).decode('utf-8').replace("*", ""))
            else:
                raise ValueError("Invalid read mode")
            if end_ip_num >= k:
                M = mid
                if mid == 0:
                    break
                high = mid - 1
            else:
                low = mid + 1
        return M

    def ip_to_int(self, ip: str) -> int:
        _ip = socket.inet_aton(ip)
        return struct.unpack("!L", _ip)[0]

    def unpack_int_4byte(self, offset: int) -> int:
        if self.read_mode == "mem":
            return struct.unpack('<L', self.data[offset:offset + 4])[0]
        elif self.read_mode == "file":
            self._handle.seek(offset)
            record = self._handle.read(4)
            return struct.unpack("<L", record[0:4])[0]
        else:
            print("Invalid read mode")
            exit()

    def unpack_int_1byte(self, offset: int) -> int:
        if self.read_mode == "mem":
            return struct.unpack('<B', self.data[offset:offset + 1])[0]
        elif self.read_mode == "file":
            self._handle.seek(offset)
            record = self._handle.read(1)
            return struct.unpack("<B", record[0:1])[0]
        else:
            print("Invalid read mode")
            exit()

    def unpack_int_8byte(self, offset: int) -> int:
        if self.read_mode == "mem":
            return struct.unpack('<Q', self.data[offset:offset + 8])[0]
        elif self.read_mode == "file":
            self._handle.seek(offset)
            record = self._handle.read(8)
            return struct.unpack("<Q", record[0:8])[0]
        else:
            print("Invalid read mode")
            exit()

    def unpack_int_2byte(self, offset: int) -> int:
        if self.read_mode == "mem":
            return struct.unpack('<H', self.data[offset:offset + 2])[0]
        elif self.read_mode == "file":
            self._handle.seek(offset)
            record = self._handle.read(2)
            return struct.unpack("<H", record[0:2])[0]
        else:
            print("Invalid read mode")
            exit()

    def get_ipv4_addr(self, j: int) -> str:
        p = 2052 + (j * self.buchang)
        offset = 0
        length = 0
        offset = self.unpack_int_4byte(4 + p)
        length = self.unpack_int_1byte(8 + p)
    
        if self.read_mode == "mem":
            return self.data[offset:offset + length].decode('utf-8')
        elif self.read_mode == "file":
            self._handle.seek(offset)
            return self._handle.read(length).decode('utf-8')
        else:
            print("Invalid read mode")
            exit()

    def get_ipv6_addr(self, j):
        p = self.numbers * 12 + 4 + 4 + (j * 55)  # numbers*12+4+4 第一段 第二段的长度       55 是 第三段 压缩 50 + 4 + 1 50是ip encode
        offset = self.unpack_int_4byte(p + 50)  # 第四部分 的 起始位置
        length = self.unpack_int_1byte(50 + p + 4)  # 第四部分的长度
        if self.read_mode == "mem":
            return self.data[offset:offset + length].decode('utf-8')
        elif self.read_mode == "file":
            self._handle.seek(offset)
            return self._handle.read(length).decode('utf-8')
        else:
            raise ValueError("Invalid read mode")


def download_db_file(down_url: str, local_file_path: str):
    """
    从下载链接下载离线库文件
    :param down_url: 离线库下载文件地址
    :param local_file_path: 离线库存放文件地址,包含文件名,必须与原文件存放位置相同,否则监视线程无法检测到文件的更改并重新加载
    """
    try:
        with requests.get(down_url, stream=True) as r:
            r.raise_for_status()
            # 从path中获取文件名
            path = urlparse(r.url).path
            file_name = str.split(path, "/")[-1]
            print(local_file_path, file_name)
            save_path = os.path.join(local_file_path, file_name)
            # 分块下载
            with open(save_path, 'wb') as f:
                for chunk in r.iter_content(chunk_size=8192):
                    f.write(chunk)
    except requests.exceptions.HTTPError as http_err:
        print(f'HTTP error occurred: {http_err}')
        return False
    except Exception as err:
        print(f'Other error occurred: {err}')
        return False
    print("下载文件成功：", save_path)
    # 如果是zip文件，解压到当前文件夹
    if save_path.endswith(".zip"):
        extract_to_dir = os.path.dirname(save_path)

        # 解压ZIP文件
        with zipfile.ZipFile(save_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to_dir)
            print(f'Extracted all files to {extract_to_dir}')

        # 删除压缩包
        os.remove(save_path)
        print(f'Removed the zip file: {save_path}')
    return True
